﻿namespace Jsl.FxCop
{
    using System;
    using System.Collections.Generic;
    using Microsoft.FxCop.Sdk;

    public class MethodVariableUsage : BinaryReadOnlyVisitor, IEnumerable<KeyValuePair<Variable, VariableUsage>>
    {
        private readonly Method method;
        private IDictionary<Variable, VariableUsage> variableUsages;

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the <see cref="MethodVariableUsage"/> class.
        /// </summary>
        /// <param name="method">The method to process.</param>
        public MethodVariableUsage(Method method)
        {
            this.method = method;
        }
        #endregion Constructors

        #region Public Properties
        /// <summary>
        /// Gets the method.
        /// </summary>
        /// <value>The method.</value>
        public Method Method
        {
            get { return this.method; }
        }

        /// <summary>
        /// Gets the <see cref="Jsl.FxCop.VariableUsage"/> for the specified variable.
        /// </summary>
        /// <value>The <see cref="Jsl.FxCop.VariableUsage"/> for the specified variable.</value>
        public VariableUsage this[Variable variable]
        {
            get
            {
                this.EnsureVariableUsageHasBeenBuilt();

                return this.variableUsages[variable];
            }
        }
        #endregion Public Properties

        #region Public Methods
        /// <summary>
        /// Gets the enumerator.
        /// </summary>
        /// <returns>The enumerator.</returns>
        public IEnumerator<KeyValuePair<Variable, VariableUsage>> GetEnumerator()
        {
            this.EnsureVariableUsageHasBeenBuilt();

            return this.variableUsages.GetEnumerator();
        }

        /// <summary>
        /// Visits the assignment statement.
        /// </summary>
        /// <param name="assignment">The assignment.</param>
        public override void VisitAssignmentStatement(AssignmentStatement assignment)
        {
            // Ignore the Pop instruction.
            if (assignment.Source.NodeType != NodeType.Pop)
            {
                // Assigning null to the variable doesn't tell us anything so ignore that as well.
                if (!assignment.Source.IsNullLiteral())
                {
                    this.UpdateVariableUsage(assignment.Target as Variable, assignment.Source.Type);
                }
            }

            base.VisitAssignmentStatement(assignment);
        }

        /// <summary>
        /// Visits the method call.
        /// </summary>
        /// <param name="call">The call.</param>
        public override void VisitMethodCall(MethodCall call)
        {
            // Get any out or ref parameters that are variables. These are considered assignments.
            var calledMethod = call.GetMethod();
            for (int i = 0; i < call.Operands.Count; i++)
            {
                (call.Operands[i] as UnaryExpression)
                    .If(unary => unary.NodeType == NodeType.AddressOf)
                    .IfNotNull(unary => unary.Operand as Variable)
                    .DoIfNotNull(variable =>
                    {
                        var reference = calledMethod.Parameters[i].Type as Reference;
                        if (reference != null)
                        {
                            this.UpdateVariableUsage(variable, reference.ElementType);
                        }
                    });
            }

            base.VisitMethodCall(call);
        }
        #endregion Public Methods

        #region Private Methods
        private void EnsureVariableUsageHasBeenBuilt()
        {
            if (this.variableUsages == null)
            {
                // Create an empty collection and process the method to fill it.
                this.variableUsages = new Dictionary<Variable, VariableUsage>();
                this.Visit(this.Method);
            }
        }
        
        private void UpdateVariableUsage(Variable variable, TypeNode type)
        {
            // Bail out here instead of checking before the call.
            if (variable == null)
            {
                return;
            }

            // Create a new VariableUsage if one doesn't already exist.
            VariableUsage variableUsage;
            if (!this.variableUsages.TryGetValue(variable, out variableUsage))
            {
                variableUsage = new VariableUsage();
                this.variableUsages.Add(variable, variableUsage);
            }

            // Add this to the variable's usage information.
            variableUsage.AddUsage(type);
        }
        #endregion Private Methods

        #region System.Collections.IEnumerator
        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return this.GetEnumerator();
        }
        #endregion System.Collections.IEnumerator
    }
}